//
//  KImageAutoEnhancer.m
//  Kia
//
//  Created by Andrey on 28/05/2009.
//  Copyright 2009 Karma Software. All rights reserved.
//

#import "KImageAutoEnhancer.h"
#import "KAutoLevels.h"
#import "KCurves.h"
#import "KEqualizeChannel.h"
#import "KFadeFilter.h"
#import "KSegmentFilter.h"
#import "KImageLayer.h"
#import "KImage.h"
#import "KCurve.h"
#import "KSimpleFloatArray.h"
#import "KSimpleUIntArray.h"
#import "KSegmentationMap.h"
#import "KImageProcessor.h"

@implementation KImageAutoEnhancer

- (id) initWithImage:(KImage*)image andSegmentationMap:(K2DSegmentationMap*)map
{
	if (self = [super init])
	{
		imageToEnhance = [image retain];
		segmentationMap = [map retain];
	}
	
	return self;
}

+ (KImageAutoEnhancer*) autoEnhancerWithImage:(KImage*)image andSegmentationMap:(K2DSegmentationMap*)map
{
	return [[[KImageAutoEnhancer alloc] initWithImage:image andSegmentationMap:map] autorelease];
}

#define OptimizationStep 0.25
#define CorrectionChannel BrightnessImageChannel
#define DEFAULT_CONTRAST_INCREASE 0.15

- (KSegmentFilter*) getDefaultFilterOfType: (enum KFilterType)filterType forSegment: (NSUInteger)segment
{
	KFilter* defaultFilter = [KImageProcessor getDefaultFilterOfType:filterType];
	
	KSegmentFilter* result = [KSegmentFilter filterWithParameters:
							  [NSDictionary dictionaryWithObjectsAndKeys:
							   defaultFilter,
							   @"Filter", 
							   segmentationMap.segments[segment],
							   @"Segment", nil
							   ]];
	
	return result;
}

- (KSegmentFilter*) constructCurvesFilterForSegment: (NSUInteger)segment
{	
	KCurves* curves;
	CGFloat lowerControlPoint = 1e30;
	CGFloat upperControlPoint = -1e30;
	
	ForEachImagePixel(imageToEnhance)
	{
		NSUInteger index = y * imageToEnhance.width + x;
		
		if (segmentationMap.crispSegmentationMap.data[index] == segment)
		{
			KRedundantColorVector colorVector = [[imageToEnhance pixelAtX:x y:y] redundantColorVector];
			
			if (colorVector.data[CorrectionChannel] < lowerControlPoint)
				lowerControlPoint = colorVector.data[CorrectionChannel];
			
			if (colorVector.data[CorrectionChannel] > upperControlPoint)
				upperControlPoint = colorVector.data[CorrectionChannel];
		}
	}
	
	/* Dirty hacks mode ON */
	lowerControlPoint *= 1.2;
	upperControlPoint *= 0.8;
	/* Dirty hacks mode OFF */
	
	CGFloat outputForLowerControlPoint = MAX(lowerControlPoint - DEFAULT_CONTRAST_INCREASE, 0.0);
	CGFloat outputForUpperControlPoint = MIN(upperControlPoint + DEFAULT_CONTRAST_INCREASE, 1.0);
	
	KCurve* curve = [[[KCurve alloc] init] autorelease];
	[curve addControlPoint:NSMakePoint(0.0, 0.0)];
	[curve addControlPoint:NSMakePoint(1.0, 1.0)];
	[curve addControlPoint:NSMakePoint(lowerControlPoint, outputForLowerControlPoint)];
	[curve addControlPoint:NSMakePoint(upperControlPoint, outputForUpperControlPoint)];
	
	NSNumber* affectedChannel = [NSNumber numberWithUnsignedInt:CorrectionChannel];
	
	curves = [KCurves filterWithParameters:
			   [NSDictionary dictionaryWithObjectsAndKeys:
				curve, 
				@"Curve", 
				affectedChannel, 
				@"AffectedChannel", nil
				]];
	
	KSegmentFilter* segmentCurves = [KSegmentFilter filterWithParameters:
									 [NSDictionary dictionaryWithObjectsAndKeys:
									  curves, 
									  @"Filter", 
									  segmentationMap.segments[segment],
									  @"Segment", nil
									  ]];
	
	return segmentCurves;
}

- (CGFloat) findOptimalMultiplierForFilter: (KFilter*)filter
{
	CGFloat multiplier = 0.0;
	CGFloat contrast = [imageToEnhance contrast];
	CGFloat uniformity = [imageToEnhance uniformity];
	CGFloat previousContrast = contrast;
	CGFloat previousUniformity = uniformity;
	
	while (multiplier <= 1.0)
	{
		previousContrast = contrast;
		previousUniformity = uniformity;
		
		NSLog(@"Multiplier: %.4f", multiplier);
		Boolean cleanUpNeeded = [KImageProcessor processImage:imageToEnhance usingTemporaryLayerForFilter:filter withMultiplier:multiplier];
		contrast = [imageToEnhance contrast];
		uniformity = [imageToEnhance uniformity];
		
		if (cleanUpNeeded)
		{
			[imageToEnhance removeLayerFromTop];		
			[imageToEnhance recomputeStatistics];
		}
		
		if (previousContrast > contrast || previousUniformity > uniformity)
			break;
		
		multiplier += OptimizationStep;
	}
	
	multiplier -= OptimizationStep;
	NSLog(@"Optimal multiplier: %.4f", multiplier);
	
	return multiplier;
}

- (KSimpleFloatArray*) findOptimalMultipliersForFilterChainForSegment: (NSUInteger)segment
{
	KSimpleFloatArray* optimalMultipliers = [KSimpleFloatArray arrayWithCapacity:2];
	
	NSLog(@"Optimizing Curves...");
	KFilter* curves;
	if (segment != -1)
		curves = [self constructCurvesFilterForSegment:segment];
	else
		curves = [KImageProcessor getDefaultFilterOfType:Curves];
	
	optimalMultipliers.data[0] = [self findOptimalMultiplierForFilter:curves];	
	Boolean cleanUpNeeded = [KImageProcessor processImage:imageToEnhance usingTemporaryLayerForFilter: curves withMultiplier:optimalMultipliers.data[0]];
	
	NSLog(@"Optimizing Equalize Channel...");
	if (segment != -1)
		optimalMultipliers.data[1] = [self findOptimalMultiplierForFilter: [self getDefaultFilterOfType: Equalize forSegment: segment]];	
	else
		optimalMultipliers.data[1] = [self findOptimalMultiplierForFilter: [KImageProcessor getDefaultFilterOfType: Equalize]];	
	
	// Optimizing curves may have added one extra layer
	if (cleanUpNeeded)
		[imageToEnhance removeLayerFromTop];
	
	return optimalMultipliers;
}

- (void) applyEnhancementFilterChainForSegment:(NSUInteger)segment withMultipliers: (KSimpleFloatArray*)multipliers
{
	NSMutableArray* enhancementFilterChain = [[NSMutableArray arrayWithCapacity:[multipliers count]] retain];
	
	if (multipliers.data[0] != 0.0)
	{
		KFilter* curves;
		if (segment != -1)
			curves = [self constructCurvesFilterForSegment:segment];
		else
			curves = [KImageProcessor getDefaultFilterOfType:Curves];
		
		KFadeFilter* fadedCurves = [KFadeFilter filterWithParameters:
									[NSDictionary dictionaryWithObjectsAndKeys:
									 curves, 
									 @"Filter", 
									 [NSNumber numberWithFloat:multipliers.data[0]], 
									 @"Fraction", nil
									 ]];
		
		if (segment != -1)
		{
			KSegmentFilter* segmentCurves = [KSegmentFilter filterWithParameters:
											 [NSDictionary dictionaryWithObjectsAndKeys:
											  fadedCurves,
											  @"Filter",
											  segmentationMap.segments[segment],
											  @"Segment", nil
											  ]];
			
			[enhancementFilterChain addObject:segmentCurves];
		}
		else
			[enhancementFilterChain addObject:fadedCurves];
	}
	
	if (multipliers.data[1] != 0.0)
	{
		KFadeFilter* fadedEqualizeChannel = [KFadeFilter filterWithParameters:
											 [NSDictionary dictionaryWithObjectsAndKeys:
											  [KImageProcessor getDefaultFilterOfType: Equalize], 
											  @"Filter", 
											  [NSNumber numberWithFloat:multipliers.data[1]], 
											  @"Fraction", nil
											  ]];
		
		if (segment != -1)
		{
			KSegmentFilter* segmentEqualizeChannel = [KSegmentFilter filterWithParameters:
													  [NSDictionary dictionaryWithObjectsAndKeys:
													   fadedEqualizeChannel,
													   @"Filter",
													   segmentationMap.segments[segment],
													   @"Segment", nil
													   ]];
			
			[enhancementFilterChain addObject:segmentEqualizeChannel];
		}
		else
			[enhancementFilterChain addObject:fadedEqualizeChannel];
	}
	
	[KImageProcessor processImage:imageToEnhance usingFilterChain:enhancementFilterChain];
	[enhancementFilterChain release];
}

- (void) autoEnhanceWithImportantSegment: (NSUInteger)segment
{
	NSLog(@"Applying lossless Auto Levels...");
	[KImageProcessor processImage:imageToEnhance usingFilter:[KImageProcessor getDefaultFilterOfType:AutoLevels]];
	NSLog(@"%@", imageToEnhance);
	
	KSimpleFloatArray* optimalMultipliers = [self findOptimalMultipliersForFilterChainForSegment:segment];
	NSLog(@"Applying enhancement filter chain...");
	[self applyEnhancementFilterChainForSegment:segment withMultipliers:optimalMultipliers];
	NSLog(@"%@", imageToEnhance);
}

- (void) dealloc
{
	[segmentationMap release];
	[imageToEnhance release];
	[super dealloc];
}

@end
